/**
 * 
 */
package edu.ou.cs.youming.agents;

import java.util.AbstractMap.SimpleEntry;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Map.Entry;
import java.util.PriorityQueue;
import java.util.Set;

import edu.ou.cs.youming.jaxb.types.Attribute;
import edu.ou.cs.youming.jaxb.types.DTNode;
import edu.ou.cs.youming.jaxb.types.Dictionary;
import edu.ou.cs.youming.jaxb.types.Feature;
import edu.ou.cs.youming.jaxb.types.Label;
import edu.ou.cs.youming.jaxb.types.attributes.WSAttribute;
import edu.ou.cs.youming.services.Services;

/**
 * @author Youming Lin
 * 
 */
public final class DTAgent implements Agent {
	private final Set<Attribute> attributes = new HashSet<Attribute>();
	private final Label defaultLabel;
	private DTNode<?> dTree;

	public DTAgent(final Label defaultLabel) {
		if (defaultLabel == null) {
			throw new NullPointerException("default label");
		}

		this.defaultLabel = defaultLabel;
	}

	@Override
	public final void train(final Map<Feature, Label> features, final Dictionary dict) {
		if (features == null) {
			throw new NullPointerException("features");
		}

		if (dict == null) {
			throw new NullPointerException("dictionary");
		}

		/*
		 * for (final String s : dict.sentimentValues.keySet()) { attributes.add(new WPAttribute(s)); }
		 */

		for (final Entry<String, Double> e : dict.sentimentValues.entrySet()) {
			attributes.add(new WSAttribute(e.getKey(), e.getValue(), 5d, 5));
		}

		dTree = generateDecisionTree(features, attributes, Label.POSITIVE);
	}

	/*
	 * (non-Javadoc)
	 * 
	 * @see edu.ou.cs.youming.agents.Agent#predict(edu.ou.cs.youming.jaxb.types.Feature)
	 */
	@Override
	public final Label predict(final Feature feature) {
		if (feature == null) {
			throw new NullPointerException("feature");
		}

		DTNode<?> currentNode = dTree;

		while (true) {
			if (currentNode.value instanceof Attribute) {
				currentNode = currentNode.children.get(((Attribute) currentNode.value)
						.featureIndex(feature));
			}
			else if (currentNode.value instanceof Label) {
				return (Label) currentNode.value;
			}
			else {
				throw new IllegalArgumentException("unsupported DTNode type: "
						+ currentNode.value.getClass());
			}
		}
	}

	/**
	 * Generates a decision tree with highest info gain attributes on top.
	 * 
	 * @param features
	 * @param attributes
	 * @param defaultLabel
	 * @return
	 */
	private static final DTNode<?> generateDecisionTree(
			final Map<Feature, Label> features, final Set<Attribute> attributes,
			final Label defaultLabel) {
		if (features == null) {
			throw new NullPointerException("examples");
		}

		if (attributes == null) {
			throw new NullPointerException("examples");
		}

		if (defaultLabel == null) {
			throw new NullPointerException("default label");
		}

		// /////////////////////////////////////////////////////////////////////
		// default case 1
		// no more examples left
		if (features.isEmpty()) {
			return new DTNode<Label>(defaultLabel);
		}

		final Map<Label, Integer> labelCount = Services.countLabels(features);

		// //////////////////////////////////////////////////////////////////////
		// default case 2
		// all remaining features have the same label
		if (labelCount.size() == 1) {
			return new DTNode<Label>(labelCount.keySet().iterator().next());
		}

		// find the most common label (modal) in remaining features
		int max = 0;
		Label modal = null;

		for (final Entry<Label, Integer> entry : labelCount.entrySet()) {
			if (entry.getValue() > max) {
				modal = entry.getKey();
				max = entry.getValue();
			}
		}

		// //////////////////////////////////////////////////////////////////////
		// default case 3
		// no more attributes left
		if (attributes.isEmpty()) {
			return new DTNode<Label>(modal);
		}

		// ///////////////////////////////////////////////////////////////////////
		// recursive case
		// find best attribute, generate subtrees
		final PriorityQueue<Entry<Attribute, Double>> bestAttributes = new PriorityQueue<Entry<Attribute, Double>>(
				attributes.size(), new Comparator<Entry<Attribute, Double>>() {
					@Override
					public int compare(final Entry<Attribute, Double> o1,
							final Entry<Attribute, Double> o2) {
						return o1.getValue().compareTo(o2.getValue());
					}

				});

		// iterate over all available attributes, add to priority queue
		for (final Attribute a : attributes) {
			final ArrayList<Map<Feature, Label>> children = new ArrayList<Map<Feature, Label>>();

			for (int i = 0; i < a.numberOfIntervals(); ++i) {
				children.add(new HashMap<Feature, Label>());
			}

			// assign each feature to child node
			for (final Entry<Feature, Label> h : features.entrySet()) {
				children.get(a.featureIndex(h.getKey())).put(h.getKey(), h.getValue());
			}

			bestAttributes
					.add(new SimpleEntry<Attribute, Double>(a, remainder(children)));
		}

		// create new node with best attribute
		final Attribute best = bestAttributes.poll().getKey();
		final HashSet<Attribute> remainingAttributes = new HashSet<Attribute>(attributes);
		remainingAttributes.remove(best);
		final DTNode<Attribute> node = new DTNode<Attribute>(best);

		// separate features by best attribute
		final ArrayList<Map<Feature, Label>> children = new ArrayList<Map<Feature, Label>>();

		for (int i = 0; i < best.numberOfIntervals(); ++i) {
			children.add(new HashMap<Feature, Label>());
		}

		for (final Entry<Feature, Label> h : features.entrySet()) {
			children.get(best.featureIndex(h.getKey())).put(h.getKey(), h.getValue());
		}

		// generate subtrees rooted at best attribute
		for (final Map<Feature, Label> m : children) {
			node.addChild(generateDecisionTree(m, remainingAttributes, modal));
		}

		return node;
	}

	/**
	 * remainder calculates the remainder after questioning using the formula:<br>
	 * remainder(label) = sum (i = 1 to v) | ((pi + ni) / (p + n) * I(pi / (pi + ni), ni / (pi + ni)))
	 * 
	 * @param features
	 * @return
	 */
	private static final double remainder(final Iterable<Map<Feature, Label>> features) {
		if (features == null) {
			throw new NullPointerException("feature");
		}

		double remainder = 0d;
		int totalSize = 0;

		// find total number of remaining features at node
		for (final Map<Feature, Label> m : features) {
			totalSize += (m == null) ? 0 : m.size();
		}

		for (final Map<Feature, Label> m : features) {
			if (m != null && !m.isEmpty()) {
				final Map<Label, Integer> labelCount = Services.countLabels(m);
				final ArrayList<Double> fractions = new ArrayList<Double>(
						labelCount.size());

				// calculate proportion of labels for each split
				for (final int i : labelCount.values()) {
					fractions.add((double) i / (double) m.size());
				}

				remainder += (double) m.size() / totalSize * information(fractions);
			}
		}

		return remainder;
	}

	/**
	 * information calculates the information bits for an Iterable of fractions using the following formula:<br>
	 * I(p(v1), p(v2), ..., p(vn)) = sum (i = i to n) | (-p(vi) * (log(p(vi))) / log(2))
	 * 
	 * @param fractions
	 * @return
	 */
	private static final double information(final Iterable<Double> fractions) {
		if (fractions == null) {
			throw new NullPointerException("fractions");
		}

		double sum = 0d;
		double information = 0d;

		for (final double d : fractions) {
			if (d < 0d) {
				throw new IllegalArgumentException("negative probability");
			}

			sum += d;
			information -= d * Math.log(d) / Math.log(2);
		}

		if (sum != 1d) {
			throw new IllegalArgumentException("probabilities do not sum to 1");
		}

		return information;
	}
}